package com.pactera.madp.common.crypto.sm;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.HexUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.SecureUtil;

import java.security.SecureRandom;

public class SM4Utils {
    /* 使用CBC模式，需要一个向量iv，可增加加密算法的强度 */
//    private static final String PARAMETER_SPEC = "1234567890123456";

	/**
	 * 数据加密[EBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @return				加密后的数据[二进制]
	 */
    public static byte[] encrypt(byte[] sourceBytes, byte[] keyBytes){
    	return encryptECBInner(sourceBytes, keyBytes);
	}

	/**
	 * 数据加密[EBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @return				加密后的数据[十六进制字符串]
	 */
	public static String encryptHex(byte[] sourceBytes, byte[] keyBytes){
		return HexUtil.encodeHexStr(encrypt(sourceBytes, keyBytes));
	}

	/**
	 * 数据加密[EBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @return				加密后的数据[base64格式]
	 */
	public static String encryptBase64(byte[] sourceBytes, byte[] keyBytes){
		return Base64.encode((encrypt(sourceBytes, keyBytes)));
	}

	/**
	 * 数据加密[EBC]模式
	 * @param source		待加密字符串
	 * @param key			密钥字符串[base64格式或十六进制格式]
	 * @return				加密后的数据[十六进制字符串]
	 */
	public static String encryptHex(String source, String key){
		return encryptHex(StrUtil.bytes(source), SecureUtil.decode(key));
	}

	/**
	 * 数据加密[EBC]模式
	 * @param source		待加密字符串
	 * @param key			密钥字符串[base64格式或十六进制格式]
	 * @return				加密后的数据[base64格式]
	 */
	public static String encryptBase64(String source, String key){
		return encryptBase64(StrUtil.bytes(source), SecureUtil.decode(key));
	}

	/**
	 * 数据加密[CBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @param ivBytes		初始向量，（字符串长度16位，转换为二进制）
	 * @return				加密后数据[二进制]
	 */
	public static byte[] encrypt(byte[] sourceBytes, byte[] keyBytes, byte[] ivBytes){
    	return encryptCBCInner(sourceBytes, keyBytes, ivBytes);
	}

	/**
	 * 数据加密[CBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @param ivBytes		初始向量，（字符串长度16位，转换为二进制）
	 * @return				加密后数据[十六进制字符串]
	 */
	public static String encryptHex(byte[] sourceBytes, byte[] keyBytes, byte[] ivBytes){
		return HexUtil.encodeHexStr(encrypt(sourceBytes, keyBytes, ivBytes));
	}

	/**
	 * 数据加密[CBC]模式
	 * @param sourceBytes	待加密二进制数据
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @param ivBytes		初始向量，（字符串长度16位，转换为二进制）
	 * @return				加密后数据[十六进制字符串]
	 */
	public static String encryptBase64(byte[] sourceBytes, byte[] keyBytes, byte[] ivBytes){
		return Base64.encode(encrypt(sourceBytes, keyBytes, ivBytes));
	}

	/**
	 * 数据加密[CBC]模式
	 * @param source		待加密字符串
	 * @param key			密钥[base64格式或十六进制格式]
	 * @param iv			初始向量[base64格式或十六进制格式]
	 * @return				加密后数据[十六进制字符串]
	 */
	public static String encryptHex(String source, String key, String iv){
		return encryptHex(StrUtil.bytes(source), SecureUtil.decode(key), SecureUtil.decode(iv));
	}

	/**
	 * 数据加密[CBC]模式
	 * @param source		待加密字符串
	 * @param key			密钥[base64格式或十六进制格式]
	 * @param iv			初始向量[base64格式或十六进制格式]
	 * @return				加密后数据[base64格式]
	 */
	public static String encryptBase64(String source, String key, String iv){
		return encryptBase64(StrUtil.bytes(source), SecureUtil.decode(key), SecureUtil.decode(iv));
	}

	/**
	 * 数据解密[EBC]模式
	 * @param cipherBytes	待解密数据[二进制]
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @return				解密后数据[二进制]
	 */
	public static byte[] decrypt(byte[] cipherBytes, byte[] keyBytes) {
		return decryptECBInner(cipherBytes, keyBytes);
	}

	/**
	 * 数据解密[EBC]模式
	 * @param cipher		待解密数据[base64格式或十六进制格式]
	 * @param key			密钥[base64格式或十六进制格式]
	 * @return				解密后数据[二进制]
	 */
	public static byte[] decrypt(String cipher, String key) {
		return decrypt(SecureUtil.decode(cipher), SecureUtil.decode(key));
	}

	/**
	 * 数据解密[BC]模式
	 * @param cipher		待解密数据[base64格式或十六进制格式]
	 * @param key			密钥[base64格式或十六进制格式]
	 * @return				解密后数据字符串
	 */
	public static String decryptStr(String cipher, String key) {
		return new String(decrypt(cipher, key));
	}


	/**
	 * 数据解密[CBC]模式
	 * @param cipherBytes	待解密数据[二进制]
	 * @param keyBytes		二进制密钥，长度 = 16
	 * @param ivBytes		初始向量，（字符串长度16位，转换为二进制）
	 * @return				解密后数据[二进制]
	 */
	public static byte[] decrypt(byte[] cipherBytes, byte[] keyBytes, byte[] ivBytes) {
		return decryptCBCInner(cipherBytes, keyBytes, ivBytes);
	}

	/**
	 * 数据解密[CBC]模式
	 * @param cipher		待解密数据[base64格式或十六进制格式]
	 * @param key			密钥[base64格式或十六进制格式]
	 * @param iv			初始向量，（字符串长度16位）
	 * @return				解密后数据[二进制]
	 */
	public static byte[] decrypt(String cipher, String key, String iv) {
		return decrypt(SecureUtil.decode(cipher), SecureUtil.decode(key), SecureUtil.decode(iv));
	}

	/**
	 * 数据解密[CBC]模式
	 * @param cipher		待解密数据[base64格式或十六进制格式]
	 * @param key			密钥[base64格式或十六进制格式]
	 * @param iv			初始向量，（字符串长度16位）
	 * @return				解密后数据[字符串]
	 */
	public static String decryptStr(String cipher, String key, String iv) {
		return new String(decrypt(cipher, key, iv));
	}

    /**
     * 随机生成一个 SM4 秘钥
     *
     * @return 字符流，长度为16
     */
    public static byte[] createSM4Key() {
        return createSM4KeyInner();
    }
    
    /* ------------------------ 内部实现 ----------------------- */
    
    /**
     * 随机生成一个 SM4 秘钥
     *
     * @return byte数组，长度为16
     */
    private static byte[] createSM4KeyInner() {
        byte[] keyBytes = new byte[16];
    
        SecureRandom random = new SecureRandom();
        random.nextBytes(keyBytes);
        return keyBytes;
    }
    
    /**
     * CBC 模式，加密
     *
     * @param srcBytes 原始数据
     * @param keyBytes 秘钥，长度 = 16
     * @param ivBytes  偏移量，长度 = 16
     * @return 加密后的数据
     */
    private static byte[] encryptCBCInner(byte[] srcBytes, byte[] keyBytes, byte[] ivBytes) {
        try {
            SM4 sm4 = new SM4();
            long[] secretKey = sm4.sm4_setkey_enc(keyBytes);
            return sm4.sm4_crypt_cbc(SM4.SM4_ENCRYPT, secretKey, true, ivBytes, srcBytes);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
    
    /**
     * CBC 模式，解密
     *
     * @param encryptedBytes 加密后的数据
     * @param keyBytes       秘钥，长度 = 16
     * @param ivBytes        偏移量，长度 = 16
     * @return 解密后的数据
     */
    private static byte[] decryptCBCInner(byte[] encryptedBytes, byte[] keyBytes, byte[] ivBytes) {
        try {
            SM4 sm4 = new SM4();
            long[] secretKey = sm4.sm4_setkey_dec(keyBytes);
            return sm4.sm4_crypt_cbc(SM4.SM4_DECRYPT, secretKey, true, ivBytes, encryptedBytes);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
    
    /**
     * ECB 模式，加密
     *
     * @param srcBytes 原始数据
     * @param keyBytes 秘钥，长度 = 16
     * @return 加密后的数据
     */
    private static byte[] encryptECBInner(byte[] srcBytes, byte[] keyBytes) {
        try {
            SM4 sm4 = new SM4();
            long[] secretKey = sm4.sm4_setkey_enc(keyBytes);
            return sm4.sm4_crypt_ecb(SM4.SM4_ENCRYPT, secretKey, true, srcBytes);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
    
    /**
     * ECB 模式，解密
     *
     * @param encryptBytes 加密后的数据
     * @param keyBytes     秘钥，长度 = 16
     * @return 解密后的数据
     */
    private static byte[] decryptECBInner(byte[] encryptBytes, byte[] keyBytes) {
        try {
            SM4 sm4 = new SM4();
            long[] secretKey = sm4.sm4_setkey_dec(keyBytes);
            return sm4.sm4_crypt_ecb(SM4.SM4_DECRYPT, secretKey, true, encryptBytes);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
}
